#include "udp.h"
#include "ip.h"
#include "icmp.h"

/**
 * @brief udp处理程序表
 * 
 */
map_t udp_table;
uint8_t LOCAL_IP[4] =NET_IF_IP;
/**
 * @brief udp伪校验和计算
 * 
 * @param buf 要计算的包
 * @param src_ip 源ip地址
 * @param dst_ip 目的ip地址
 * @return uint16_t 伪校验和
 */
static uint16_t udp_checksum(buf_t *buf, uint8_t *src_ip, uint8_t *dst_ip)
{
    // 12字节的伪头部
    buf_add_header(buf,sizeof(udp_peso_hdr_t));
    // 做头部的备份
    uint8_t store[sizeof(udp_peso_hdr_t)];
    memcpy(store,buf->data,sizeof(udp_peso_hdr_t));
    udp_peso_hdr_t * fake_head = (udp_peso_hdr_t *) buf->data;
    memcpy(fake_head->src_ip,src_ip,4);
    memcpy(fake_head->dst_ip,dst_ip,4);
    fake_head->placeholder=0;
    // 这里可以这么写吗？既然是udp_ckecksum,默认的ip层的上层协议是UDP
    fake_head->protocol = NET_PROTOCOL_UDP; 
    // 不算伪头部的长度
    fake_head->total_len16 = swap16(buf->len-sizeof(udp_peso_hdr_t));
    // 错误点：奇数个字节的填充放在checksum里面进行
    uint16_t csum = checksum16((uint16_t *)buf->data,buf->len);
    memcpy(buf->data,store,sizeof(udp_peso_hdr_t));
    buf_remove_header(buf,sizeof(udp_peso_hdr_t));
    return csum;
}

/**
 * @brief 处理一个收到的udp数据包
 * 
 * @param buf 要处理的包
 * @param src_ip 源ip地址
 */
void udp_in(buf_t *buf, uint8_t *src_ip)
{
    if (buf->len<sizeof(udp_hdr_t))
    {
        return;
    }
    udp_hdr_t *header = (udp_hdr_t*)buf->data;
    if (buf->len<swap16(header->total_len16))
    {
        return;
    }
    uint16_t buf_csum = header->checksum16;
    header->checksum16 = 0;
    uint16_t csum = udp_checksum(buf,src_ip,LOCAL_IP);
    if (csum != buf_csum)
    {
        return;
    }
    uint16_t des_port = swap16(header->dst_port16);
    udp_handler_t *handler = map_get(&udp_table,&des_port);
    // 如果找到处理函数
    if (handler)
    {
        buf_remove_header(buf,sizeof(udp_hdr_t));
        uint16_t src_port =  swap16(header->src_port16);
        (*handler)(buf->data, buf->len,src_ip,src_port);
        return;
    }
    // 如果找不到处理函数
    else{
        // ip数据报头部应该原本就存在，把长度加回去就好
        buf_add_header(buf,sizeof(ip_hdr_t));
        icmp_unreachable(buf,src_ip,ICMP_CODE_PORT_UNREACH);
    }
}

/**
 * @brief 处理一个要发送的数据包
 * 
 * @param buf 要处理的包
 * @param src_port 源端口号
 * @param dst_ip 目的ip地址
 * @param dst_port 目的端口号
 */
void udp_out(buf_t *buf, uint16_t src_port, uint8_t *dst_ip, uint16_t dst_port)
{
    buf_add_header(buf,sizeof(udp_hdr_t));
    udp_hdr_t *header =(udp_hdr_t *) buf->data;
    header->src_port16 = swap16(src_port);
    header->dst_port16 = swap16(dst_port);
    header->checksum16 = 0;
    header->total_len16 = swap16(buf->len);
    uint16_t csum = udp_checksum(buf,LOCAL_IP,dst_ip);
    header->checksum16 =csum;
    ip_out(buf,dst_ip,NET_PROTOCOL_UDP);
}


/**
 * @brief 初始化udp协议
 * 
 */
void udp_init()
{
    map_init(&udp_table, sizeof(uint16_t), sizeof(udp_handler_t), 0, 0, NULL);
    net_add_protocol(NET_PROTOCOL_UDP, udp_in);
}

/**
 * @brief 打开一个udp端口并注册处理程序
 * 
 * @param port 端口号
 * @param handler 处理程序
 * @return int 成功为0，失败为-1
 */
int udp_open(uint16_t port, udp_handler_t handler)
{
    return map_set(&udp_table, &port, &handler);
}

/**
 * @brief 关闭一个udp端口
 * 
 * @param port 端口号
 */
void udp_close(uint16_t port)
{
    map_delete(&udp_table, &port);
}

/**
 * @brief 发送一个udp包
 * 
 * @param data 要发送的数据
 * @param len 数据长度
 * @param src_port 源端口号
 * @param dst_ip 目的ip地址
 * @param dst_port 目的端口号
 */
void udp_send(uint8_t *data, uint16_t len, uint16_t src_port, uint8_t *dst_ip, uint16_t dst_port)
{
    buf_init(&txbuf, len);
    memcpy(txbuf.data, data, len);
    udp_out(&txbuf, src_port, dst_ip, dst_port);
}